from ripser import ripser
from persim import plot_diagrams
import matplotlib.pyplot as plt
import numpy as np
from sklearn import datasets
import sys

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

# data = datasets.make_circles(n_samples=100)[0] + 5 * datasets.make_circles(n_samples=100)[0]

# 制造一个 torus 数据
n = 1500
data = np.zeros((n, 3))
for i in range(n):
    progress_bar(i, n)
    theta = np.random.rand() * 2 * np.pi
    phi = np.random.rand() * 2 * np.pi
    data[i, 0] = (3 + np.cos(theta)) * np.cos(phi)
    data[i, 1] = (3 + np.cos(theta)) * np.sin(phi)
    data[i, 2] = np.sin(theta)

# # 将这个 torus 进行一些扭曲
# for i in range(n):
#     progress_bar(i, n)
#     theta = np.random.rand() * 2 * np.pi
#     phi = np.random.rand() * 2 * np.pi
#     data[i, 0] += 0.5 * np.cos(theta)
#     data[i, 1] += 0.5 * np.sin(theta)
#     data[i, 2] += 0.5 * np.sin(phi)

# # 制造一个 solid torus 数据
# n = 1000
# data = np.zeros((n, 3))
# for i in range(n):
#     progress_bar(i, n)
#     theta = np.random.rand() * 2 * np.pi
#     phi = np.random.rand() * 2 * np.pi
#     r = np.random.rand()
#     data[i, 0] = (2 + r * np.cos(theta)) * np.cos(phi)
#     data[i, 1] = (2 + r * np.cos(theta)) * np.sin(phi)
#     data[i, 2] = r * np.sin(theta)

# 将 data 绘制到一个三维图中
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(data[:, 0], data[:, 1], data[:, 2])
plt.show()

# 计算 data 的距离矩阵
from scipy.spatial.distance import pdist, squareform
dist = squareform(pdist(data))
print(dist.shape)

# 计算 persistence diagram
# dgms = ripser(data, maxdim=2, coeff=47)['dgms']
dgms = ripser(dist, maxdim=2, coeff=47, distance_matrix=True)['dgms']

for dgm in dgms:
    print(dgm.shape)
plot_diagrams(dgms, show=True, lifetime=True)
